import pandas as pd
import os
import glob
import soundfile as sf
import librosa as lb

dom_df = pd.read_csv("../../src/extracted/DESED_FSD50K_domestic_sounds.csv")
lm_df = pd.read_csv("../../src/extracted/DESED_FSD50K_label_mapping_non_single_speech.csv")

all_fnames = set(dom_df["fname"]).union(set(lm_df["fname"]))

dev_folder = "/nvmework1/shaonian/Datasets/FSD50K/FSD50K.dev_audio"
eval_folder = "/nvmework1/shaonian/Datasets/FSD50K/FSD50K.eval_audio"

out_dir = "/nvmework1/shaonian/Datasets/FSD50K/extracted/"
if not os.path.exists(out_dir):
    os.makedirs(out_dir)
found = 0
for f in glob.glob(os.path.join(dev_folder, "*")) + glob.glob(os.path.join(eval_folder, "*")):
    fname = os.path.basename(f).replace(".wav", "")
    if int(fname) in all_fnames:
        out_path = os.path.join(out_dir, fname + ".wav")
        if not os.path.exists(out_path):
            # load. resample to 16kHz
            audio, sr = sf.read(f)
            if sr != 16000:
                audio = lb.resample(audio, orig_sr=sr, target_sr=16000)
            sf.write(out_path, audio, 16000)
            found += 1
    
print(f"Total files found: {found} out of {len(all_fnames)}")